import os
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

import numpy as np
import torch
from FlagEmbedding import BGEM3FlagModel, FlagReranker
from typing import List, Tuple
import faiss


class BGE_M3_EmbeddingRerankSystem:
  def __init__(self,
      embedding_model_path: str = "BAAI/bge-m3",
      reranker_model_path: str = "BAAI/bge-reranker-v2-m3",
      device: str = "cuda" if torch.cuda.is_available() else "cpu"):
    """
    初始化BGE-M3 Embedding和Reranker模型

    参数:
        embedding_model_path: Embedding模型路径
        reranker_model_path: Reranker模型路径
        device: 计算设备 (cuda/cpu)
    """
    self.device = device

    # 1. 加载Embedding模型
    print("⏳ 正在加载Embedding模型...")
    self.embedding_model = BGEM3FlagModel(
        embedding_model_path,
        use_fp16=True
    )

    # 2. 加载Reranker模型
    print("⏳ 正在加载Reranker模型...")
    self.reranker_model = FlagReranker(
        reranker_model_path,
        use_fp16=True
    )

    # 3. 初始化FAISS索引
    self.dim = 1024  # BGE-M3的dense向量维度
    self.faiss_index = faiss.IndexFlatIP(self.dim)

    print("✅ 系统初始化完成!")

  def index_documents(self, documents: List[str]) -> None:
    """
    为文档集合创建向量索引

    参数:
        documents: 文档列表
    """
    print(f"📚 正在为 {len(documents)} 个文档创建索引...")
    self.documents = documents
    embeddings = self.embedding_model.encode(
        documents,
        batch_size=12,
        max_length=8192
    )['dense_vecs'].astype('float32')
    faiss.normalize_L2(embeddings)
    self.faiss_index.add(embeddings)
    print("✅ 文档索引创建完成!")

  def retrieve_and_rerank(self,
      query: str,
      top_k_recall: int = 100,
      top_k_rerank: int = 5) -> List[Tuple[str, float]]:
    """
    执行召回和重排流程

    参数:
        query: 用户查询
        top_k_recall: 召回阶段返回的文档数量
        top_k_rerank: 重排后返回的文档数量

    返回:
        重排后的文档列表 (文档内容, 相关性分数)
    """
    # 阶段1: 召回 (Embedding)
    print("🔍 召回阶段...")
    query_embedding = self.embedding_model.encode(
        [query],
        batch_size=1,
        max_length=8192
    )['dense_vecs'].astype('float32').reshape(1, -1)
    faiss.normalize_L2(query_embedding)
    distances, indices = self.faiss_index.search(query_embedding, top_k_recall)
    recall_indices = indices[0]
    recall_scores = distances[0]
    recall_docs = [self.documents[idx] for idx in recall_indices if idx != -1]

    print(f"  召回完成: {len(recall_docs)} 个文档")

    # 阶段2: 重排 (Reranker)
    print("📊 重排阶段...")

    # 准备输入对
    pairs = [
      [query, doc]
      for doc in recall_docs
    ]

    # 计算重排分数
    rerank_scores = self.reranker_model.compute_score(pairs, normalize=True)

    # 组合召回分数和重排分数
    combined_scores = [
      (recall_docs[i], rerank_scores[i], recall_scores[i])
      for i in range(len(recall_docs))
    ]

    # 按重排分数排序
    combined_scores.sort(key=lambda x: x[1], reverse=True)

    # 返回Top K结果
    return [(doc, rerank_score) for doc, rerank_score, _ in
            combined_scores[:top_k_rerank]]


# ==================== 使用示例 ====================
if __name__ == "__main__":
  # 1. 初始化系统
  system = BGE_M3_EmbeddingRerankSystem(
      embedding_model_path="BAAI/bge-m3",
      reranker_model_path="BAAI/bge-reranker-v2-m3",
      device="cpu"
  )

  # 2. 准备文档库
  documents = [
    "北京是中国的首都，位于华北平原的北部。",
    "巴黎是法国的首都，也是欧洲重要的文化和商业中心。",
    "重力是物体之间相互吸引的力，由艾萨克·牛顿首次系统描述。",
    "量子力学是描述微观粒子行为的物理学理论。",
    "长城是中国古代的军事防御工程，是世界文化遗产。",
    "埃菲尔铁塔位于法国巴黎，是世界上最著名的建筑之一。",
    "黑洞是时空曲率大到连光都无法逃脱的天体。",
    "故宫位于北京，是中国明清两代的皇家宫殿。",
    "相对论是爱因斯坦提出的关于时空和引力的理论。",
    "亚马逊雨林是世界上最大的热带雨林，位于南美洲。",
    "大熊猫是中国的国宝，主要生活在四川的山区。",
    "人工智能是计算机科学的一个分支，旨在创造智能机器。",
    "长城始建于春秋战国时期，用于防御北方游牧民族。",
    "量子纠缠是量子力学中的一种现象，两个粒子无论距离多远都能瞬间影响对方。",
    "卢浮宫位于巴黎，是世界上最大的艺术博物馆之一。"
  ]

  # 3. 创建文档索引
  system.index_documents(documents)

  # 4. 执行查询
  queries = [
    "中国的首都是哪个城市？",
    "解释什么是量子纠缠",
    "巴黎有哪些著名景点？",
    "什么是相对论？"
  ]

  for query in queries:
    print(f"\n🔎 查询: '{query}'")
    results = system.retrieve_and_rerank(query, top_k_recall=5, top_k_rerank=3)

    print("📈 重排结果:")
    for i, (doc, score) in enumerate(results):
      print(
        f"  {i + 1}. [分数: {score:.4f}] {doc[:60]}{'...' if len(doc) > 60 else ''}")